from torch.utils.data import DataLoader, TensorDataset

import os
import sys

path = os.path.abspath(os.path.join(os.path.dirname(__file__), '../../../'))
sys.path.append(path)

from shell.knowledge_graph.text_classification.data_process import *
from shell.knowledge_graph.text_classification.model import *


# 加载训练数据与验证数据
def load_data_all(batch_size=32):
    df = pd.read_csv(data_path)
    train_x = df['question']
    train_y = df['label']
    train_x = padding_seq(train_x.apply(seq2index))
    train_y = np.array(train_y)

    train_data_set = TensorDataset(torch.from_numpy(train_x), torch.from_numpy(train_y))
    train_data_loader = DataLoader(dataset=train_data_set, batch_size=batch_size, shuffle=True)

    return train_data_loader, None, None


def load_data(batch_size=32):
    # df = pd.read_csv('parent_data.csv')
    df = pd.read_csv(data_path)
    train_df, eval_df = split_data(df)
    train_x = train_df['question']
    train_y = train_df['label']
    eval_x = eval_df['question']
    eval_y = eval_df['label']

    train_x = padding_seq(train_x.apply(seq2index))
    train_y = np.array(train_y)

    train_data_set = TensorDataset(torch.from_numpy(train_x), torch.from_numpy(train_y))
    train_data_loader = DataLoader(dataset=train_data_set, batch_size=batch_size, shuffle=True)

    eval_x = padding_seq(eval_x.apply(seq2index))
    return train_data_loader, eval_x, eval_y.values


# 训练模型
def train():
    train_data_loader, eval_x, eval_y = load_data(batch_size)
    model = TextCNN(n_classes)
    # model = BiLSTM(vocab_len,
    #                embedding_size,
    #                lstm_hidden_size,
    #                lstm_layer_nums,
    #                lstm_bidirectional,
    #                drop_out,
    #                n_classes)
    model.fit(train_data_loader, (eval_x, eval_y), epochs)


if __name__ == '__main__':
    train()
